package com.bz.websocket;

import com.bz.websocket.error.CommonErrorCode;
import com.bz.websocket.exception.BizException;
import com.bz.websocket.utils.AssertUtil;
import com.bz.websocket.utils.ObjectMapperUtil;
import com.fasterxml.jackson.core.type.TypeReference;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Component;
import org.springframework.web.reactive.function.client.WebClient;
import org.springframework.web.reactive.socket.WebSocketHandler;
import org.springframework.web.reactive.socket.WebSocketMessage;
import org.springframework.web.reactive.socket.WebSocketSession;

import java.util.Map;

import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

/**
 * @author zhongyongbin
 * @since 2021/5/10
 */
@Component
public class BizWebSocketHandler implements WebSocketHandler {

    @Autowired
    private AuthManager authManager;

    @Autowired
    private UserSessionManager userSessionManager;

    @Autowired
    private WebClient webClient;

    @Override
    public Mono<Void> handle(WebSocketSession session) {
        Flux<WebSocketMessage> resp = session.receive()
                .map(msg -> {
                    RequestProtocol protocol = ObjectMapperUtil.transferToObject(msg.getPayloadAsText(), RequestProtocol.class);
                    AssertUtil.requireNonNull(protocol, () -> new BizException(CommonErrorCode.PROTOCOL_ERROR));

                    return new BizProcess(session.getId(), protocol);
                })
                .flatMap(bizProcess -> {
                    String sessionId = bizProcess.getSessionId();
                    if (bizProcess.getProtocol().getAction() == ProtocolAction.CONNECT) {
                        return authManager.auth(bizProcess.getProtocol())
                                .flatMap(res -> res ? userSessionManager.addSession(sessionId, new UserSession(sessionId, bizProcess.getProtocol().getUserId())) : Mono.just(false))
                                .map(res -> {
                                    bizProcess.setAuth(res);
                                    return bizProcess;
                                })
                                .doOnError(e -> bizProcess.setAuth(false))
                                .onErrorReturn(bizProcess);
                    } else {
                        bizProcess.setAuth(true);
                    }
                    return Mono.just(bizProcess);
                })
                .doOnNext(bizProcess -> {
                    String sessionId = bizProcess.getSessionId();
                    if (bizProcess.getProtocol().getAction() == ProtocolAction.COMMUNICATION) {
                        bizProcess.setAuth(userSessionManager.refreshExpireTime(sessionId));
                    }
                })
                .flatMap(bizProcess -> {
                    if (!bizProcess.getAuth()) {
                        return Mono.just(session.textMessage(ObjectMapperUtil.transferToString(new ResponseProtocol(CommonErrorCode.SESSION_NOT_ALIVE))));
                    }

                    RequestProtocol requestProtocol = bizProcess.getProtocol();
                    BizRequest bizRequest = new BizRequest();
                    bizRequest.setBiz(requestProtocol.getBiz());
                    bizRequest.setDevice(requestProtocol.getDevice());
                    bizRequest.setUserId(requestProtocol.getUserId());
                    bizRequest.setPayload(requestProtocol.getPayload());

                    return webClient.post()
                            .uri(WebSocketBiz.getInstance(bizProcess.getProtocol().getBiz()).getBizUrl())
                            .contentType(MediaType.APPLICATION_JSON)
                            .bodyValue(ObjectMapperUtil.transferToString(bizRequest))
                            .retrieve()
                            .toEntity(String.class)
                            .map(response -> {
                                if (response.getStatusCode().is2xxSuccessful()) {
                                    Map<String, Object> payload = null;
                                    if (response.hasBody()) {
                                        String responseString = response.getBody();
                                        payload = ObjectMapperUtil.transferToObject(responseString, new TypeReference<Map<String, Object>>() {
                                        });
                                    }

                                    ResponseProtocol responseProtocol = new ResponseProtocol();
                                    responseProtocol.setPayload(payload);

                                    return session.textMessage(ObjectMapperUtil.transferToString(requestProtocol));
                                } else {
                                    return session.textMessage(ObjectMapperUtil.transferToString(new ResponseProtocol(CommonErrorCode.SYSTEM_ERROR)));
                                }
                            })
                            .onErrorReturn(session.textMessage(ObjectMapperUtil.transferToString(new ResponseProtocol(CommonErrorCode.SYSTEM_ERROR))));
                });

        return session.send(resp).then();
    }
}
